# Use python to submit jobs automatically.
import os
import re
import subprocess
import sys
import time


def submit_and_get_job_id(qsub_command):
    # Execute the qsub command
    print(f"run cmd: {qsub_command}")
    process = subprocess.Popen(
        qsub_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
    # Get the output of the command
    stdout, stderr = process.communicate()
    # If the command fails, raise an error
    if process.returncode != 0:
        raise Exception(
            f"cmd '{qsub_command}' failed with error: {stderr.decode('utf-8', errors='ignore')}")
    # The output is similar to "108235.eic"
    output = stdout.decode('utf-8', errors='ignore').strip()
    # Extract the job ID from the output (assuming the job ID is the first number in the output)
    job_id = output.split('.')[0]
    print(f"Submitted job: {job_id} @",
          time.strftime("%Y-%m-%d %H:%M:%S", time.localtime()))
    return job_id


def count_jobs_by_queue(qstat_output, queue):
    # Use regular expressions to match the queue of jobs and count the number of jobs in the specified queue
    # Assume each line format is "JobID Name User TimeUse Status Queue"
    job_queue_pattern = r'^\d+\.\w+\s+\S+\s+\S+\s+\S+\s+\S+\s+(\S+)$'
    # Find all matching queue fields
    job_queues = re.findall(job_queue_pattern, qstat_output, re.MULTILINE)
    # Count the number of jobs in the specified queue
    cnt = job_queues.count(queue)
    print(f"Number of jobs with queue '{queue}': {cnt}")
    return cnt


def modify_source_file(file_path, data_dict):
    """
    Modify the 'xymwdc' format source params in the given file.

    :param file_path: The path to the file to be modified.
    :param data_dict: A dictionary containing the new values for x, y, z, tbeg, trise, mag, strike, dip, and rake.
    """
    # print(f"Modifying source params to: {data_dict}")
    with open(file_path, 'r') as file:
        lines = file.readlines()

    modified_lines = []
    for line in lines:
        if line.strip().startswith("#") or line.strip() == "":
            # Skip comment lines or blank lines
            modified_lines.append(line)
        else:
            # Split the line into columns based on spaces
            columns = line.split()

            # Check if the line has the correct number of columns for 'xymwdc' format
            if len(columns) == 9:  # x, y, z, tbeg, trise, mag, strike, dip, rake
                # Replace the values in the line with the corresponding values from data_dict
                new_line = f"{data_dict['x']:6.1f} {data_dict['y']:5.1f} {data_dict['z']:5.1f} " \
                           f"{data_dict['tbeg']:5.2f} {data_dict['trise']:6.3f} {data_dict['mag']:6.1f} " \
                           f"{data_dict['strike']:8.1f} {data_dict['dip']:6.1f} {data_dict['rake']:6.1f}\n"
                modified_lines.append(new_line)
                print(f"Modified source from:{line} to: {new_line.strip()}")
            else:
                # Otherwise, just keep the original line
                modified_lines.append(line)

    # Write the modified lines back to the file
    with open(file_path, 'w') as file:
        file.writelines(modified_lines)


if __name__ == "__main__":
    source_path = './chengyu/source_chengyu.dat'
    joblist_path = './chengyu/joblist_chengyu_250430.dat'
    job_path = 'job_fcy_C'
    result_path = './out_chengyu/wav'
    start_index = 304
    print(f"Start from job {start_index}.")
    sys.stdout.flush()
    # Read the job file, execute each job, and record the last completed job to autosubmit_log.txt
    with open(joblist_path, 'r') as file:
        lines = file.readlines()
    for index, line in enumerate(lines):
        if index < start_index:
            print(f"Skip job {index}.")
            continue
        print(f"[job {index}].")
        sys.stdout.flush()
        params = line.strip().split()
        # Modify the source file
        data_dict = {
            'x': float(params[0]),
            'y': float(params[1]),
            'z': float(params[2]),
            'tbeg': 0.1,
            'trise': 0.11,
            'mag': 1.0,
            'strike': float(params[3]),
            'dip': float(params[4]),
            'rake': float(params[5])
        }
        modify_source_file(source_path, data_dict)
        # Submit the job
        job_id = submit_and_get_job_id(f"qsub {job_path}")
        # Poll the job status
        while True:
            sys.stdout.flush()
            time.sleep(60)
            qstat_output = subprocess.check_output(
                "qstat", stderr=subprocess.STDOUT, shell=True).decode('utf-8')
            if job_id not in qstat_output:
                print(f"Job {job_id} has finished.")
                break
            else:
                # check if output file exists
                if not os.path.exists(f'{job_path}.e{job_id}'):
                    print(f"Job {job_id} still queuing.")
                    continue
                # job is running, check output
                with open(f'{job_path}.e{job_id}', 'r') as file:
                    output = file.readlines()
                print(
                    f"Job {job_id} is running. output line: {len(output)}")
        # check job output and rename output dir
        with open(f'{job_path}.e{job_id}', 'r') as file:
            output = file.read()
        if 'Total time' in output:
            print(f"Job {job_id} has successfully completed.")
            # rename output dir
            os.rename(result_path, f'{result_path}_{index}')
            print(f"Rename {result_path} to {result_path}_{index}")
        else:
            print(f"Job {job_id} has failed.")
            break
